# encoding: utf-8
# !/usr/local/bin/python3

"""
@author: Gao Shuo
@contact: dorothy400@163.com
@software: PyCharm
@file: roof_clip.py
@time: 2020/4/21 11:48

前期探索：
    卫星图片大小不同
    处理后的标签是18，其他是0

input:
    1. 原图像文件夹
    2. 标注灰度图像文件夹（label是18）
    3. 输出黑白标注-剪切后文件夹
    4. 输出原图像-剪切后文件夹
    5. 输出01标注-剪切后文件夹
    6. 是否包含标注。如果是新的卫星图则不包含标注

output:
    剪切后的各个图像
"""

import os.path
import sys

from PIL import Image


def new_folder (path):
    """
    如果路径不是-，不存在的话就新建
    """
    if path != '-':
        if not os.path.exists(path):
            os.makedirs(path)
    return path

if __name__ == "__main__":
    path_ori = sys.argv[1]
    path_label = sys.argv[2]
    dest_255_clip = new_folder(sys.argv[3])
    dest_ori_clip = new_folder(sys.argv[4])
    dest_1_clip = new_folder(sys.argv[5])
    if_label = sys.argv[6]
    li = os.listdir(path_ori)
    li = [x.replace('.jpg','') for x in li]
    target_width = 800
    target_height = 480
    for filepath in li:
        file_ori = os.path.join(path_ori, filepath + '.jpg')
        im_ori = Image.open(file_ori)
        if_save = True # 是否保留所有图
        im_label = im_ori
        if if_label == '1':
            file_label = os.path.join(path_label, filepath + '.png')
            im_label = Image.open(file_label)
        else:
            if_save = True
        size = im_ori.size
        n_width = size[0] / target_width
        n_height = size[1] / target_height
        filepath = filepath.replace(" ","")
        print filepath
        print n_width
        print n_height
        for i in range(n_width):
            for j in range(n_height):
                image_ori = im_ori.crop((target_width*i, target_height*j,
                                         target_width*(i+1), target_height*(j+1)))
                if if_label == '1':
                    image_255 = Image.new("L", (target_width, target_height))
                    image_1 = Image.new("L", (target_width, target_height))
                    for my_i in range(target_width):
                        for my_j in range(target_height):
                            target_pixel = (i*target_width + my_i, j*target_height + my_j)
                            color_label = im_label.getpixel(xy=target_pixel)
                            if int(color_label) > 0:
                                if_save=True
                                image_255.putpixel((my_i, my_j), 255)
                                image_1.putpixel((my_i, my_j), 1)
                            else:
                                image_255.putpixel((my_i, my_j), 0)
                                image_1.putpixel((my_i, my_j), 0)
                    if if_save:
                        image_255.save(os.path.join(dest_255_clip, filepath + '_' + str(i) + '_' + str(j) + '.png'))
                        image_1.save(os.path.join(dest_1_clip, filepath + '_' + str(i) + '_' + str(j) + '.png'))

                if if_save:
                    image_ori.save(os.path.join(dest_ori_clip, filepath+'_'+str(i) +'_'+ str(j) + '.png'))
                    print ("saved: \n"+ filepath+'_'+str(i) +'_'+ str(j) + '.png')

        print 'finished: ', filepath

